package com.chj.leetcode.tijie450to500;

import java.util.Comparator;
import java.util.HashMap;
import java.util.Map;
import java.util.PriorityQueue;

public class Code480_SlidingWindowMedian {
	public double[] medianSlidingWindow(int[] nums, int k) {
		double[] medians = new double[nums.length - k + 1];
		Map<Integer, Integer> hashTable = new HashMap<>();
		for (int i = 0; i < nums.length; i++) {
			hashTable.put(nums[i], 0);
		}

//		PriorityQueue<Integer> maxHeap = new PriorityQueue<Integer>(new Comparator<Integer>() {
//
//			@Override
//			public int compare(Integer o1, Integer o2) {
//				return o2.compareTo(o1);
//			}
//
//		});
//
//		PriorityQueue<Integer> minHeap = new PriorityQueue<Integer>(new Comparator<Integer>() {
//
//			@Override
//			public int compare(Integer o1, Integer o2) {
//				return o1.compareTo(o2);
//			}
//
//		});

		int i = 0; // index of current incoming element being processed
//
//		// initialize the heaps
//		while (i < k) {
//			maxHeap.add(nums[i++]);
//		}
//		for (int j = 0; j < k / 2; j++) {
//			minHeap.add(maxHeap.peek());
//			maxHeap.poll();
//		}

		while (i < k) {
			addNum(nums[i++]);
		}

		int index = 0;
		while (true) {
			// get median of current window
//			medians[index++] = (k & 1) > 0 ? maxHeap.peek() : ((double) maxHeap.peek() + (double) minHeap.peek()) * 0.5;
			medians[index++] = findMedian();
			if (i >= nums.length) {
				break; // break if all elements processed
			}

			int out_num = nums[i - k]; // outgoing element
			int in_num = nums[i++]; // incoming element
//			int balance = 0; // balance factor

			// number `out_num` exits window
//			balance += (out_num <= maxHeap.peek() ? -1 : 1);
			hashTable.put(out_num, hashTable.get(out_num) + 1);
			addNum(in_num);

			// number `in_num` enters window
//			if (!maxHeap.isEmpty() && in_num <= maxHeap.peek()) {
//				balance++;
//				maxHeap.add(in_num);
//			} else {
//				balance--;
//				minHeap.add(in_num);
//			}

			// re-balance heaps
//			if (balance < 0) { // `maxHeap` needs more valid elements
//				maxHeap.add(minHeap.peek());
//				minHeap.poll();
//				balance++;
//			}
//			if (balance > 0) { // `minHeap` needs more valid elements
//				minHeap.add(maxHeap.peek());
//				maxHeap.poll();
//				balance--;
//			}

			// remove invalid numbers that should be discarded from heap tops
			while (!maxHeap.isEmpty() && hashTable.containsKey(maxHeap.peek()) && hashTable.get(maxHeap.peek()) > 0) {
				int removeNum = maxHeap.poll();
				hashTable.put(removeNum, hashTable.get(removeNum) - 1);
			}
			while (!minHeap.isEmpty() && hashTable.containsKey(minHeap.peek()) && hashTable.get(minHeap.peek()) > 0) {
				int removeNum = minHeap.poll();
				hashTable.put(removeNum, hashTable.get(removeNum) - 1);
			}
		}

		return medians;
	}

	private PriorityQueue<Integer> maxHeap = new PriorityQueue<Integer>(new MaxHeapComparator());
	private PriorityQueue<Integer> minHeap = new PriorityQueue<Integer>(new MinHeapComparator());

	private void modifyTwoHeapsSize() {
		if (this.maxHeap.size() == this.minHeap.size() + 2) {
			this.minHeap.add(this.maxHeap.poll());
		}
		if (this.minHeap.size() == this.maxHeap.size() + 2) {
			this.maxHeap.add(this.minHeap.poll());
		}
	}

	public void addNum(int num) {
		if (maxHeap.isEmpty() || num <= maxHeap.peek()) {
			maxHeap.add(num);

		} else {
			minHeap.add(num);
		}
		modifyTwoHeapsSize();
	}

	public double findMedian() {
		int maxHeapSize = this.maxHeap.size();
		int minHeapSize = this.minHeap.size();
		if (maxHeapSize + minHeapSize == 0) {
			return 0.0;
		}
		Integer maxHeapHead = this.maxHeap.peek();
		Integer minHeapHead = this.minHeap.peek();
		if (maxHeapSize == minHeapSize) {
			return (maxHeapHead + minHeapHead) * 0.5;
		}
		return maxHeapSize > minHeapSize ? maxHeapHead : minHeapHead;
	}

	public class MaxHeapComparator implements Comparator<Integer> {
		@Override
		public int compare(Integer o1, Integer o2) {
			if (o2 > o1) {
				return 1;
			} else {
				return -1;
			}
		}
	}

	public class MinHeapComparator implements Comparator<Integer> {
		@Override
		public int compare(Integer o1, Integer o2) {
			if (o2 < o1) {
				return 1;
			} else {
				return -1;
			}
		}
	}
}
